package com.zhouchj.elasticsearch.basic;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.text.StrPool;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.reflect.TypeToken;
import com.zhouchj.elasticsearch.constant.EsConstant;
import com.zhouchj.elasticsearch.model.page.PageForm;
import com.zhouchj.elasticsearch.model.page.QueryPageEntity;
import com.zhouchj.elasticsearch.model.page.Tuple;
import org.apache.commons.lang3.StringUtils;
import org.apache.lucene.search.join.ScoreMode;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.AggregatorFactories;
import org.elasticsearch.search.aggregations.bucket.nested.Nested;
import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedTerms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.TopHitsAggregationBuilder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Sort;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;

import java.lang.reflect.Type;
import java.time.Instant;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * ES工具类
 *
 * @author zhouchj
 * @date 2023/7/21
 */
public class EsQueryBuilderUtil {

    /**
     * 每个桶的数量上限 暂时设置为1000
     */
    private static final Integer AGG_MAX_SIZE = EsConstant.AGG_MAX_SIZE;

    private static final Gson ES_GSON = new GsonBuilder()
            .disableHtmlEscaping()
            .serializeSpecialFloatingPointValues()
            .setFieldNamingStrategy(new EsFieldNamingPolicy())
            .create();

    private EsQueryBuilderUtil() {
    }

    /**
     * 构造分页器
     *
     * @param queryBuilder 查询构造器
     * @param pageForm     分页入参
     */
    public static void buildPage(NativeSearchQueryBuilder queryBuilder, PageForm pageForm) {
        buildPage(queryBuilder, pageForm, Sort.unsorted());
    }


    /**
     * 构造分页器
     *
     * @param queryBuilder 查询构造器
     * @param pageForm     分页入参
     * @param sort         排序
     */
    public static void buildPage(NativeSearchQueryBuilder queryBuilder, PageForm pageForm, Sort sort) {
        long current = pageForm.getCurrent() - 1;
        long size = pageForm.getSize();
        if (current < 0) {
            current = 0;
        }
        PageRequest page = PageRequest.of((int) current, (int) size, sort);
        queryBuilder.withPageable(page);
    }

    /**
     * 构造查询
     *
     * @param builder   最顶层的条件构造器
     * @param fieldName 字段名
     * @param value     字段值 字段值之间关系为or
     */
    public static void buildTermQuery(BoolQueryBuilder builder, String fieldName, List<?> value) {
        if (CollUtil.isEmpty(value)) {
            return;
        }
        BoolQueryBuilder internalQueryBuild = getTermBoolQueryBuilder(fieldName, value);
        if (Objects.nonNull(internalQueryBuild)) {
            builder.must(internalQueryBuild);
        }
    }

    /**
     * 构造非查询
     *
     * @param builder   最顶层的条件构造器
     * @param fieldName 字段名
     * @param value     字段值 字段值之间关系为or
     */
    public static void buildNotTermQuery(BoolQueryBuilder builder, String fieldName, List<?> value) {
        if (CollUtil.isEmpty(value)) {
            return;
        }
        BoolQueryBuilder internalQueryBuild = getTermBoolQueryBuilder(fieldName, value);
        if (Objects.nonNull(internalQueryBuild)) {
            builder.mustNot(internalQueryBuild);
        }
    }

    /**
     * 获取BoolQueryBuilder
     *
     * @param fieldName 字段名称
     * @param value     值
     * @return BoolQueryBuilder
     */
    private static BoolQueryBuilder getTermBoolQueryBuilder(String fieldName, List<?> value) {
        BoolQueryBuilder internalQueryBuild = QueryBuilders.boolQuery();
        // 过滤空数据
        List<?> collect = value.stream()
                .filter(EsQueryBuilderUtil::filterNullValue)
                .collect(Collectors.toList());
        if (collect.isEmpty()) {
            return null;
        }
        internalQueryBuild.should(QueryBuilders.termsQuery(fieldName, collect));
        return internalQueryBuild;
    }


    private static boolean filterNullValue(Object r) {
        if (Objects.isNull(r)) {
            return false;
        }
        if (r instanceof String && StringUtils.isBlank((String) r)) {
            return false;
        }
        return true;
    }

    /**
     * 构造嵌套查询
     *
     * @param builder         最顶层的条件构造器
     * @param nestedFieldName 嵌套字段名称
     * @param fieldName       字段名
     * @param value           字段值 字段值之间关系为or
     */
    public static void buildNestedTermQuery(BoolQueryBuilder builder, String nestedFieldName, String fieldName, List<?> value) {
        if (CollUtil.isEmpty(value)) {
            return;
        }
        BoolQueryBuilder internalQueryBuild = QueryBuilders.boolQuery();
        for (Object o : value) {
            if (Objects.isNull(o)) {
                continue;
            }
            if (o instanceof String && StringUtils.isBlank((String) o)) {
                continue;
            }
            QueryBuilder subFieldQuery = QueryBuilders.termQuery(nestedFieldName + StrPool.DOT + fieldName, o);
            internalQueryBuild.should(QueryBuilders.nestedQuery(nestedFieldName, subFieldQuery, ScoreMode.None));
        }
        builder.must(internalQueryBuild);
    }

    /**
     * 构造嵌套查询
     *
     * @param builder         最顶层的条件构造器
     * @param nestedFieldName 嵌套字段名称
     * @param fieldName       字段名
     * @param value           字段值 字段值之间关系为or
     */
    public static void buildNestedTermQuery(BoolQueryBuilder builder,
                                            String nestedFieldName,
                                            String fieldName,
                                            Object value) {
        buildNestedTermQuery(builder, nestedFieldName, fieldName, Collections.singletonList(value));
    }

    /**
     * 构造查询
     *
     * @param builder   最顶层的条件构造器
     * @param fieldName 字段名
     * @param value     字段值 字段值之间关系为or
     */
    public static void buildTermQuery(BoolQueryBuilder builder, String fieldName, Object value) {
        EsQueryBuilderUtil.buildTermQuery(builder, fieldName, Collections.singletonList(value));
    }


    /**
     * 构造范围查询
     *
     * @param builder   最顶层的条件构造器
     * @param fieldName 字段名
     * @param startTime 开始时间
     * @param endTime   结束时间
     */
    public static void buildRangeQuery(BoolQueryBuilder builder, String fieldName, Instant startTime, Instant endTime) {
        if (Objects.isNull(startTime) || Objects.isNull(endTime)) {
            return;
        }
        if (startTime.isAfter(endTime)) {
            throw new RuntimeException("开始时间不能晚于结束时间");
        }
        RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery(fieldName).gte(startTime).lte(endTime);
        builder.must(rangeQueryBuilder);
    }


    /**
     * 构造分组
     *
     * @param builder 查询构造器
     */
    public static void buildTermsAgg(NativeSearchQueryBuilder builder, String bucketName, String field) {
        TermsAggregationBuilder agg = AggregationBuilders
                .terms(bucketName)
                .field(field)
                .size(AGG_MAX_SIZE);
        builder.addAggregation(agg);
    }


    /**
     * 分组构造器
     *
     * @param builder         查询构造器
     * @param bucketName      桶名(返回中分组的名称)
     * @param nestedFieldName 嵌套字段名称
     * @param subFieldName    子字段名称
     */
    public static void buildNestedAgg(NativeSearchQueryBuilder builder,
                                      String bucketName,
                                      String nestedFieldName,
                                      String subFieldName) {
        TermsAggregationBuilder subAgg = AggregationBuilders.terms(bucketName)
                .field(nestedFieldName + StrPool.DOT + subFieldName)
                .size(AGG_MAX_SIZE);
        NestedAggregationBuilder agg = AggregationBuilders.nested(bucketName, nestedFieldName)
                .subAggregation(subAgg);
        builder.addAggregation(agg);
    }

    /**
     * 获取字段分组后前topSize条数据
     *
     * @param builder        查询构造器
     * @param bucketName     外层分组桶名
     * @param topBucketName  内存top桶名
     * @param topSize        每个分组中的数据条数
     * @param fieldName      需要分组的字段
     * @param orderFieldName 排序字段（可为null）
     */
    public static void buildTopAgg(NativeSearchQueryBuilder builder,
                                   String bucketName,
                                   String topBucketName,
                                   int topSize,
                                   String fieldName,
                                   String orderFieldName) {
        TermsAggregationBuilder aggregationBuilder = AggregationBuilders.terms(bucketName)
                .field(fieldName)
                .size(AGG_MAX_SIZE);

        TopHitsAggregationBuilder topHit = AggregationBuilders
                .topHits(topBucketName)
                .size(topSize)
                .fetchField(fieldName);

        if (StringUtils.isNotBlank(orderFieldName)) {
            topHit.sort(orderFieldName);
        }
        aggregationBuilder.subAggregations(AggregatorFactories.builder().addAggregator(topHit));
        builder.addAggregation(aggregationBuilder);
    }


    /**
     * 构造PageEntity返回
     *
     * @param search   查询返回
     * @param pageForm 入参中的分页参数
     * @return 构造PageEntity返回
     */
    public static <T> QueryPageEntity<T> buildPageEntity(SearchHits<T> search, PageForm pageForm) {
        QueryPageEntity<T> pageEntity = new QueryPageEntity<>();
        List<T> collect = search.getSearchHits()
                .stream()
                .map(SearchHit::getContent)
                .collect(Collectors.toList());
        pageEntity.setRecords(collect);
        pageEntity.setTotal(search.getTotalHits());
        pageEntity.setCurrent(pageForm.getCurrent());
        pageEntity.setSize(pageForm.getSize());

        addConditionToPage(search.getAggregations(), pageEntity);
        return pageEntity;
    }

    /**
     * 往pageEntity中加入查询条件
     *
     * @param aggregations 聚合信息
     * @param pageEntity   分页信息
     */
    private static <T> void addConditionToPage(Aggregations aggregations, QueryPageEntity<T> pageEntity) {
        processBucket(aggregations,
                name -> true,
                (name, bucket) -> pageEntity.addCondition(name, bucket.getKeyAsString(),
                        bucket.getKeyAsString(),
                        (int) bucket.getDocCount()));
    }

    /**
     * 查询桶下的所有数量
     *
     * @param aggregations 分组
     * @param bucketName   桶名
     * @return Tuple
     */
    public static List<Tuple> getBucketInfo(Aggregations aggregations, String bucketName) {
        List<Tuple> result = new ArrayList<>();
        processBucket(aggregations,
                name -> name.equals(bucketName),
                (name, bucket) -> {
                    Tuple tuple = Tuple.of(bucket.getKeyAsString(), bucket.getKeyAsString(), (int) bucket.getDocCount());
                    result.add(tuple);
                });
        return result;
    }

    /**
     * 处理分组
     *
     * @param aggregations 分组
     * @param filterBucket 过滤桶 返回true的不过滤 如bucketName->"a".equals(bucketName);
     * @param doSomething  桶中的每个字段需要做什么
     */
    public static void processBucket(Aggregations aggregations,
                                     Function<String, Boolean> filterBucket,
                                     BiConsumer<String, Terms.Bucket> doSomething) {
        if (Objects.isNull(aggregations)) {
            return;
        }
        List<Aggregation> aggList = aggregations.asList();
        for (Aggregation agg : aggList) {
            if (agg instanceof ParsedTerms) {
                ParsedTerms parsedStringTerms = (ParsedTerms) agg;
                String name = agg.getName();
                Boolean apply = filterBucket.apply(name);
                if (!Boolean.TRUE.equals(apply)) {
                    continue;
                }
                List<? extends Terms.Bucket> buckets = parsedStringTerms.getBuckets();
                for (Terms.Bucket bucket : buckets) {
                    doSomething.accept(name, bucket);
                }
            } else if (agg instanceof Nested) {
                Nested nested = (Nested) agg;
                processBucket(nested.getAggregations(), filterBucket, doSomething);
            } else {
                // 未定义
                throw new UnsupportedOperationException("Aggregation未解析");
            }

        }
    }


    /**
     * 把Es返回出来的String转换成实体类
     *
     * @return 实体类
     */
    public static <T> T convertStringToEsEntity(String json, Class<T> clazz) {
        try {
            return ES_GSON.fromJson(json, clazz);
        } catch (Exception e) {
            throw new RuntimeException("解析失败", e);
        }
    }

    /**
     * 把Es返回出来的String转换成实体类
     *
     * @return 实体类
     */
    public static <T> String convertEntityToString(T obj) {
        try {
            return ES_GSON.toJson(obj);
        } catch (Exception e) {
            throw new RuntimeException("解析失败", e);
        }
    }


    /**
     * 把Es返回出来的String转换成Map
     *
     * @return 实体类
     */
    public static <T> List<Map<String, Object>> convertEntityToListMap(T obj) {
        try {
            Type mapType = new TypeToken<List<Map<String, Object>>>() {
            }.getType();
            String str = ES_GSON.toJson(obj);
            return ES_GSON.fromJson(str, mapType);
        } catch (Exception e) {
            throw new RuntimeException("解析失败", e);
        }
    }


    /**
     * 把Es返回出来的String转换成Map
     *
     * @return 实体类
     */
    public static <T> Map<String, Object> convertEntityToMap(T obj) {
        try {
            Type mapType = new TypeToken<Map<String, Object>>() {
            }.getType();
            String str = ES_GSON.toJson(obj);
            return ES_GSON.fromJson(str, mapType);
        } catch (Exception e) {
            throw new RuntimeException("解析失败", e);
        }
    }

}
